### Functions to generate grouped forest plots
### Version: 0.4
### Date: 20150914
### Authors: cl14022

### This script does:
### 1. Read in a set of meta-analysis data that  arise from mr_singlesnp
### 2. Order sub-analyses alphabetically
### 2.1. Order the studies within each sub-analysis by effect size
### 3. Generate a list of spacings so that sub-analyses can have headers and be spaced apart for easy reading
### 4. Generate a properly-spaced forest plot (without a summary effect) for all studies
### 5. Generate a properly-spaced set of annotations for the left-hand side (lhs) of the figure
### 6. Generate a properly-spaced set of annotations for the right-hand side (rhs) of the figure
### 7. Combine the three sub-figures into a single figure
### 7.1.Generate a list of spacings for annotation columns and for the forest-plot column
### 8. Export results as a pdf/windows meta file and to screen

# library(ggplot2)
# library(plyr)
# library(gtable)

## 8.0.1. Unifiying all into a single function that calls each of the above methods

#' Grouped forest plot
#'
#' @param name (character) name of the delimited file containing all of the results on the first sheet (needs to have headers), or of the r object.
#' @param eff_Col (character) name of the column in the delimited file that contains the effect sizes.
#' @param exposure_Name (character) name of the column in the delimited file containing the *types* of studies.
#' @param outcome_Name (character) name of the column in the delimited file containing the names of each study.
#' @param outfile_Name (character) name to be used for output file (*.pdf) or (*.wmf).
#' @param forest_Title (character) the title to be used for a forest plot.
#' @param left_Col_Names (character vector) vector containing the names of the left-hand-side annotation columns in the delimited file.
#' @param left_Col_Titles (character vector) vector containing the titles for each left-hand-side annotation column.
#' @param right_Col_Names (character vector) vector containing the names of the right-hand-side annotation columns in the delimited file.
#' @param right_Col_Titles (character vector) vector containing the titles for each right-hand-side annotation column.
#' @param debug (logical) show warnings `TRUE`/`FALSE`?
#' @param log_ES (logical) perform natural log transform of effect sizes and confidence bounds `TRUE`/`FALSE`?
#' @param decrease (logical) sort the studies by decreasing effect sizes `TRUE`/`FALSE`?
#' @param se_Col (character) name of the column giving the standard error of the effect sizes.
#' @param returnRobj (logical) return the graph as an internal R object `TRUE`/`FALSE`?
#'
#' @keywords internal
#' @return grid object giving the forest plot (or plot as pdf)
mr_forest_plot_grouped <-
  function(
    name,
    eff_Col = "b",
    exposure_Name = "exposure",
    outcome_Name = "outcome",
    forest_Title = '',
    outfile_Name = 'annot_FP.pdf',
    left_Col_Names = c("Exposure", "Outcome"),
    left_Col_Titles = NULL,
    right_Col_Names = c("p", "Outcome.n.case", "Outcome.n.control", "Outcome.sample.size"),
    right_Col_Titles = NULL,
    debug = FALSE,
    log_ES = FALSE,
    decrease = TRUE,
    returnRobj = TRUE,
    se_Col = "se"
  ) {
    # name = (character) name of the r object from mr_singlesnp()
    # inRobj = (logical) is the data to be used an internal R object y/n?
    # eff_Col = (character) name of the column in the delimited file that contains the effect sizes

    # exposure_Name = (character) name of the column in the delimited file containing the exposure to use
    # outcome_Name = (character) name of the column in the delimited file containing outcome to be analyzed
    # outfile_Name = (character) name to be used for output file (*.pdf) or (*.wmf)
    # forest_Title = (character) the title to be used for a forest plot
    # left_Col_Names = (character vector) vector containing the names of the left-hand-side annotation columns in the delimited file
    # left_Col_Titles = (character vector) vector containing the titles for each left-hand-side annotation column
    # right_Col_Names = (character vector) vector containing the names of the right-hand-side annotation columns in the delimited file
    # right_Col_Titles = (character vector) vector containing the titles for each right-hand-side annotation column
    # debug = (logical) show warnings y/n?

    # log_ES = (logical) perform natural log transform of effect sizes and confidence bounds y/n?
    # decrease = (logical) sort the studies by decreasing effect sizes y/n?

    # returnRobj = (logical) return the graph as an internal R object y/n?
    name$lb <- name[, eff_Col] - stats::qnorm(p = 0.95) * name[, se_Col]
    name$ub <- name[, eff_Col] + stats::qnorm(p = 0.95) * name[, se_Col]
    data <- name
    spacer <- function(exposure, eff_col, outcome, Data_Fm, decrease = TRUE) {
      # exposure = (character), column name of the column containing the name of the particular exposure
      # eff_col = (character), column name of the column giving the effect size of interest
      # outcome = (character), column name giving study names/phenotype names
      # Data_Fm = (name) name of the data frame containing the meta-analytic cata
      # decrease = (logical) order the studies by decreasing effect size y/n?

      N_study <- nrow(Data_Fm)

      # Count the number of categories of studies (e.g. by type of illness )
      N_sub <- length(unique(Data_Fm[, exposure]))

      # The layout of rows is one per study, then, buffering room added for the subgroups
      N_space <- N_study + 2 * N_sub

      spacing_vec <- vector(mode = "numeric", length = N_space)
      content_list <- list()
      attr_list <- list()
      row_list <- list()
      idx <- 1
      for (subgp in unique(Data_Fm[, exposure])) {
        # loop over substudy ids to fill the 'spacing vector,' which determines placement of the rows in the forest plot
        sg_idxs <- which(Data_Fm[, exposure] == subgp)
        n_std <- length(sg_idxs)
        spacing_vec[[idx]] <- idx
        spacing_vec[[(idx + 1 + n_std)]] <- (idx + 1 + n_std)
        std_sort <- order(as.numeric(Data_Fm[sg_idxs, eff_col]), decreasing = decrease)
        spacing_vec[((idx + 1):(idx + n_std))] <- idx + 1:n_std

        # In the same loop, fill a list with study names but also with headers to serve as the primary annotation of the forest plot
        content_list[[idx]] <- subgp
        content_list[[(idx + 1 + n_std)]] <- ""
        content_list[((idx + 1):(idx + n_std))] <- Data_Fm[sg_idxs[std_sort], outcome]

        # A list of attributes for the typeface so that headers are in bold face
        attr_list[[idx]] <- "bold"
        attr_list[[(idx + 1 + n_std)]] <- "plain"
        attr_list[((idx + 1):(idx + n_std))] <- "plain"

        # The most important list, contains mapping between spaces on the forest plot and rows in the meta-analytic data frame
        row_list[[idx]] <- NA
        row_list[[(idx + 1 + n_std)]] <- NA
        row_list[((idx + 1):(idx + n_std))] <- sg_idxs[std_sort]

        idx <- idx + 2 + n_std
      }
      # returns data frame giving spacings, primary annotation (content_list), typeface attributes for the primary annotation (attr_list), and  mapping between rows in forest plot and rows in meta-analytic data frame
      return(data.frame(
        spacing_vec = (1 + length(spacing_vec) - spacing_vec),
        content_list = unlist(content_list),
        row_list = unlist(row_list),
        attr_list = unlist(attr_list)
      ))
    }

    space_Out <- function(data_Fm, space_Fm) {
      # a function to expand the initial data frame and add the spacing vector in a proper way
      # data_Fm = (name) data frame containing the meta-analytic data
      # space_Fm = (name) data frame containing the spacing vector and the mapping between spaces in the forest plot and rows in the meta-analytic data
      exp_data_Fm <- data.frame(space_Fm)
      for (jj in colnames(data_Fm)) {
        exp_data_Fm[, jj] <- data_Fm[space_Fm$row_list, jj]
      }

      # Returns a data frame containing spacing, primary annotations, typeface attributes for the primary annotation, and correctly-mapped rows in the meta-analytic data frame
      return(exp_data_Fm)
    }

    ### 4. Generate a properly-spaced forest plot (without a summary effect) for all studies
    ggforest <- function(
      data_Fm,
      space_col,
      eff_col,
      lb_col = "lb",
      ub_col = "ub",
      title_text = '',
      log_ES = FALSE
    ) {
      # data_Fm = (name) spaced data frame (must be from output from space_Out()) containing meta-analytic data and appropriate annotation, spacing cols etc
      # space_col = (character) column name giving the column listing spacing of effects in the forest plot
      # eff_col = (character) column name of effect-size column
      # lb_col = (character) column name of the confidence interval lower bound column
      # ub_col = (character) column name of the confidence interval upper bound column
      # title_text = (character) title text for main part of forest plot (will use column name if no text is supplied)
      # log_ES = (logical) convert effect sizes to natural log scale y/n?

      if (nchar(title_text) <= 1) {
        title_text <- eff_col
      }
      ## reassignment of column names to avoid ggplot scope problems
      data_Fm$space_col <- data_Fm[, space_col]
      data_Fm$lb_col <- data_Fm[, lb_col]
      if (log_ES) {
        data_Fm$lb_col <- log(as.numeric(data_Fm[, lb_col]))
      }
      data_Fm$ub_col <- data_Fm[, ub_col]
      if (log_ES) {
        data_Fm$ub_col <- log(as.numeric(data_Fm[, ub_col]))
      }
      data_Fm$eff_col <- data_Fm[, eff_col]
      if (log_ES) {
        data_Fm$eff_col <- log(as.numeric(data_Fm[, eff_col]))
      }
      # ggplot code to generate the forest plot using geom_segments and geom_points and to make a relatively minimal theme
      raw_forest <- ggplot2::ggplot(
        data = data_Fm,
        ggplot2::aes(
          y = space_col,
          yend = space_col,
          x = as.numeric(lb_col),
          xend = as.numeric(ub_col)
        )
      ) +
        ggplot2::geom_segment() +
        ggplot2::geom_point(ggplot2::aes(y = space_col, x = as.numeric(eff_col), size = 4)) +
        ggplot2::theme_bw() +
        ggplot2::theme(
          axis.text.y = ggplot2::element_blank(),
          axis.ticks.y = ggplot2::element_blank(),
          axis.title = ggplot2::element_blank(),
          panel.grid = ggplot2::element_blank(),
          rect = ggplot2::element_blank(),
          title = ggplot2::element_text(size = 23),
          legend.position = 'none'
        ) +
        ggplot2::expand_limits(y = c(data_Fm[, space_col] - 1, data_Fm[, space_col] + 2)) +
        ggplot2::labs(title = title_text) # returns ggplot2 object with the (un-annotated) forest plot
      return(raw_forest)
    }

    ### 5. Generate a properly-spaced set of annotations for the left-hand side (lhs) of the figure
    ### and
    ### 6. Generate a properly-spaced set of annotations for the right-hand side (rhs) of the figure
    ## 5.1. Function to generate a single column of annotations
    anot_col <- function(data_Fm, text_col, space_col, title_text = '') {
      # data_Fm = (name) space_Out data frame
      # text_col = (character) name of column containing annotations
      # title_text = (character) title of annotation column, defaults to the column name if nothing is entered

      data_Fm$space_col <- data_Fm[, space_col]
      data_Fm$text_col <- data_Fm[, text_col]

      # A hard rule to set the width of the annotation column, which sometimes truncates very wide columns (complex disease names, numbers with 16 digits, etc)
      text_widths <- c(-1, max(10, 0.5 * max(sapply(as.character(data_Fm[, text_col]), nchar))))

      # GGplot rendering of the annotation column
      lefttext <- ggplot2::ggplot(
        data = data_Fm,
        ggplot2::aes(y = space_col, x = 0, label = text_col, fontface = attr_list)
      ) +
        ggplot2::geom_text(hjust = 0) +
        ggplot2::theme_bw() +
        ggplot2::theme(
          axis.text.y = ggplot2::element_blank(),
          axis.ticks.y = ggplot2::element_blank(),
          axis.text.x = ggplot2::element_text(colour = "white"),
          axis.ticks.x = ggplot2::element_line(colour = "white"),
          axis.title = ggplot2::element_blank(),
          rect = ggplot2::element_blank(),
          panel.grid = ggplot2::element_blank(),
          title = ggplot2::element_text(size = 23)
        ) +
        ggplot2::expand_limits(
          x = text_widths,
          y = c(data_Fm[, space_col] - 1, data_Fm[, space_col] + 2)
        ) +
        ggplot2::labs(title = title_text, size = 40) # returns two-item list with left_text, the GGplot annotations, and text_widths, the x-axis limits of the plot
      return(list(left_text = lefttext, text_widths = text_widths))
    }

    ## 5.2. Function to aggregate all of the left-hand (all of the right-hand) notation columns into a single subplot, which works by applying the anot_col function to each item in the list of annotation columns
    anot_side <- function(data_Fm, space_col, col_names, title_list = '') {
      # data_Fm = (name) space_Out data frame containing effect sizes, all annotations, and spacings
      # space_col = (character) name of the column containing the spacing vector
      # col_names = (character vector) vector of the column names for the columns to be the left-hand-side or right-hand-side annotations
      # title_list = (character vector) vector of the titles to be given to each of the annotation columns, defaults to the column names

      relative_widths <- vector(mode = "numeric", length = length(col_names))
      output <- vector(mode = "list")

      if (length(title_list) <= 1) {
        title_list <- col_names
      }

      for (i in seq_along(col_names)) {
        # loop to get the widths of each annotation column and to group the annotation objects together
        col <- anot_col(
          data_Fm = data_Fm,
          text_col = col_names[i],
          space_col = space_col,
          title_text = title_list[[i]]
        )
        relative_widths[i] <- col$text_widths[2] - col$text_widths[1]
        output[[col_names[i]]] <- ggplotGrob(col$left_text)
      }
      # calculate the widths of the annotation columns, relative to each other, and then multiply it by a third so that it fits onto one side of the plot
      output$relative_widths <- relative_widths / (sum(relative_widths)) * 0.33
      # returns an output object containing: each annotation plot, listed under its own name, and the relative_widths of the columns
      return(output)
    }

    ### 7. Combine the three sub-figures into a single figure
    ### 7.1.Generate a list of spacings for annotation columns and for
    group_Plots <- function(forst_Pt, left_Hs, right_Hs) {
      # forst_Pt = (name) the name of the forest plot object
      # left_Hs = (name) the name of the aggregate annotation object for the left hand side of the plot
      # right_Hs  = (name) the name of the aggregate annotation object for the right hand side of the plot

      # Aggregate all of the plots into a single grid object for plotting
      grob_Bag <- vector(mode = 'list')

      left_RW <- left_Hs$relative_widths
      left_Grobs <- left_Hs
      left_Grobs$relative_widths <- NULL

      for (i in seq_along(left_Grobs)) {
        grob_Bag[paste('l', names(left_Grobs)[i], sep = '')] <- left_Grobs[i]
      }

      grob_Bag$m_forest <- ggplotGrob(forst_Pt)

      right_RW <- right_Hs$relative_widths
      right_Grobs <- right_Hs
      right_Grobs$relative_widths <- NULL
      for (i in seq_along(right_Grobs)) {
        grob_Bag[paste('r', names(right_Grobs)[i], sep = '')] <- right_Grobs[i]
      }

      width_vec <- c(left_RW, 0.34, right_RW)
      width_vec <- width_vec / sum(width_vec)
      # convert the grid objects (now grouped) into a table of grid objects that can be plotted using grid.draw
      grp_FP <- gtable::gtable_matrix(
        name = "groupplot",
        grobs = matrix(grob_Bag, nrow = 1),
        widths = unit(width_vec, "npc"),
        heights = unit(1, "npc")
      )

      # return the grid object table, to be plotted
      return(grp_FP)
    }

    ## act on debug flag
    if (debug == FALSE) {
      options(warn = -1)
    }

    ## read in data: beginning of the OVERALL FUNCTION

    ## order and structure effect sizes and CIs for forest plot
    space1 <- spacer(
      exposure = exposure_Name,
      eff_col = eff_Col,
      outcome = outcome_Name,
      Data_Fm = data,
      decrease = decrease
    )

    expand_data <- space_Out(data_Fm = data, space_Fm = space1)

    ## Make the forest plot
    fo1 <- ggforest(
      data_Fm = expand_data,
      space_col = 'spacing_vec',
      eff_col = eff_Col,
      lb_col = "lb",
      ub_col = "ub",
      log_ES = log_ES,
      title_text = forest_Title
    )

    ## Construct left-hand-side annotations
    left <- anot_side(
      data_Fm = expand_data,
      space_col = 'spacing_vec',
      col_names = left_Col_Names,
      title_list = left_Col_Titles
    )

    ## Construct right-hand-side annotations
    right <- anot_side(
      data_Fm = expand_data,
      space_col = 'spacing_vec',
      col_names = right_Col_Names,
      title_list = right_Col_Titles
    )

    ## group all plots together
    group <- group_Plots(forst_Pt = fo1, left_Hs = left, right_Hs = right)
    ## draw and export the annotated forest plot
    grid.newpage()
    grid.draw(group)
    if (returnRobj == FALSE) {
      grDevices::pdf(outfile_Name, width = 23.4, height = 16.5)
      grid.draw(group)
      grDevices::dev.off()
    } else {
      return(group)
    }
    ## make sure not to leave warnings turned off
    options(warn = 0)
  }
